# -*- coding: utf-8 -*-
# @Time    : 2019/2/25 12:54
# @Author  : cj
# @File    : data_item.py
import os
import cv2
import numpy as np
from utils.basic import cv_imread, cv_imwrite

from dataset.annotation import Annotation


class DataItem:
    """
    类DataItem，形如[image,[annotation1,annotation2,.],..]，为一张图片和对应的xml的信息
    """

    def __init__(self, image=None, annotation=None):
        self._image = image
        self._annotation = annotation

    def get_image(self):
        """
        :return: 返回图片
        """
        return self._image

    def set_image(self, image: np.array):
        """
        设置图片
        :param image: 更改后的图片
        """
        self._image = image

    def get_label_item(self):
        """
        :return:返回 label_item
        """
        return self._annotation

    def set_label_item(self, label_item):
        """
        设置label_item
        :param label_item: 更改后的label_item
        """
        self._annotation = label_item

    def _computer_ratio(self, max_edge, min_edge):
        """
        函数作用:在faster_rcnn 中首先对图片进行缩放，此计算图片的缩放比例
        :param max_edge: 最大边
        :param min_edge: 最小边
        :return: 缩放比例
        """
        min_shape = min(self._annotation.get_image_size()[0], self._annotation.get_image_size()[1])
        max_shape = max(self._annotation.get_image_size()[0], self._annotation.get_image_size()[1])
        min_ratio = float(min_edge) / min_shape
        max_ratio = float(max_edge) / max_shape

        return max_ratio if max_shape * min_ratio > max_edge else min_ratio

    def get_scales(self, max_edge=None, min_edge=None):
        """
        用于返回scales，返回xml里每一个boundingbox的scale，如果不设置max_edge，min_edge
        则返回原比例，否则返回图片的缩放比例下的scale
        :param max_edge: 配置文件里的最大边，用于计算图片的缩放比例
        :param min_edge: 配置文件里的最小边，用于计算图片的缩放比例
        :return: scales
        """
        scales = []
        for boundingbox in self._annotation.get_boundingboxes():
            if max_edge is None:
                scales.append(boundingbox.get_scale())
            else:
                scales.append(boundingbox.get_scale() * self._computer_ratio(max_edge, min_edge))
        return scales

    def load(self, img_path: str, xml_path: str):
        """
        加载图片的XML
        :param img_path: 图片路径
        :param xml_path: XML路径
        """
        if not os.path.exists(img_path):
            raise FileNotFoundError(img_path + ' Not Found')
        if not os.path.exists(xml_path):
            raise FileNotFoundError(xml_path + ' Not Found')
        image = cv_imread(img_path)
        # opencv读取的图片的通道为BGR
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        self._image = image
        annotation = Annotation()
        annotation.load_from_xml_file(xml_path)
        self._annotation = annotation

    def save(self, save_dir: str, databese='Unknown'):
        """
        保存图片和xml
        :param save_dir: 保存的路径
        :param databese: 数据库来源
        """
        image = self._image
        # 因为读进来的image都转换为了RGB，用CV2保存时转换成BGR
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        cv_imwrite(os.path.join(save_dir, self._annotation.get_image_filename()), image)
        # 因为读进image_path来的image都转换为了RGB，用CV2保存时转换成BGR
        self._annotation.save_to_xml_file(path=save_dir, database1=databese)

    @property
    def annotation(self):
        return self._annotation

    @property
    def image(self):
        return self._image
